import vtkmodules.all as vtk


class JsonTransformer:

    @staticmethod
    def poly_data_to_json(poly_data: vtk.vtkPolyData):
        # 获取点数据
        points = poly_data.GetPoints()
        vertices = []
        for i in range(points.GetNumberOfPoints()):
            point = points.GetPoint(i)
            vertices.append(list(point))

        # 获取面片数据
        faces = []
        polys = poly_data.GetPolys()
        polys.InitTraversal()
        idList = vtk.vtkIdList()
        while polys.GetNextCell(idList):
            face = []
            for i in range(idList.GetNumberOfIds()):
                face.append(idList.GetId(i))
            faces.append(face)

        # 构建JSON数据结构
        mesh_json = {
            'vertices': vertices,
            'faces': faces
        }

        # 获取点数据属性
        point_data = poly_data.GetPointData()
        if point_data.GetNumberOfArrays() > 0:
            attributes = {}
            for i in range(point_data.GetNumberOfArrays()):
                array = point_data.GetArray(i)
                name = array.GetName()
                data = []
                for j in range(array.GetNumberOfTuples()):
                    if array.GetNumberOfComponents() == 1:
                        data.append(array.GetValue(j))
                    else:
                        tuple_data = []
                        for k in range(array.GetNumberOfComponents()):
                            tuple_data.append(array.GetComponent(j, k))
                        data.append(tuple_data)
                attributes[name] = data
            mesh_json['point_attributes'] = attributes

        return mesh_json
